from math import nan
from sympy import Piecewise
from sympy.abc import x

def piecewise(*args):
    """
    生成分段函数或表达式。每个参数应该是一对(cond, val)，其中cond是一个布尔表达式，val是cond为真时函数的值。
    如果在输入参数中没有给出默认值，那么如果没有一个cond满足条件，会返回nan
    如果给出otherwise_val即默认值，那么当没有cond满足的时候会赋该值
    example:
    pw = piecewise(x < 0, -x, x > 0, x**2)
    """
    n = len(args)
    # without the args for ohterwise_val
    if n % 2 == 0:
        conds, vals = args[0:n:2], args[1:n:2]
        return Piecewise(*zip(vals, conds), (nan, True))
    # with
    else:
        conds, vals = args[0:n-1:2], args[1:n-1:2]
        otherwise_val = args[-1]
        return Piecewise(*zip(vals, conds), (otherwise_val, True))



pw = piecewise(x < 0, -x, x > 0, x**2)
print(pw.evalf(subs={x: -1}))   # 输出结果为 1.0
print(pw.evalf(subs={x: 0.5}))  # 输出结果为 0.25
print(pw.evalf(subs={x: 1}))    # 输出结果为 1
print(pw.evalf(subs={x: 0}))    # 输出结果为 NaN

pw2 = piecewise(x < 0, -x, x > 0, x**2,1024)
print(pw2.evalf(subs={x: 0}))   #输出结果为 1024

